'''
Created on Oct 29, 2009

@author: mkiyer
'''

from veggie.app.alignment.bowtie import Bowtie
import sys
import re
import xml.etree.cElementTree as etree
import operator

class AlignmentParams(object):
    num_mismatches = 3

class AlignmentStats(object):
    '''
    Alignment statistics measurements:
    
    insert_size: read1 + read2 + insert
    mate1_mm: # of reads with 0-mismatches, 1-mm, 2-mm, .. n-mm
    mate2_mm: # of reads with 0-mismatches, 1-mm, 2-mm, .. n-mm
    mate1_read_mm: percent mm along the read
    mate2_read_mm: percent mm along the read    
    '''
    
    def __init__(self, aln_params):
        self.aln_params = aln_params
        self.isize_mean = None
        self.isize_stdev = None
        self.isize_pdf = collections.defaultdict(lambda: 0)
        self.read_mm = None
        self.mm = np.zeros((2, aln_params.num_mismatches), dtype=np.int)
    
    def to_xml(self):
        root = etree.Element("alignment_statistics")
        isize = root.SubElement("insert_size")
        mean = isize.SubElement("mean")
        mean.text = '%d' % (int(round(self.isize_mean)))
        stdev = isize.SubElement("stdev")
        stdev.text = '%d' % (int(round(self.isize_stdev)))
        pdf = isize.SubElement("pdf")
        pdf_strings = ','.join([i for i in self.isize_pdf.items()])
        pdf.text = '|'.joing(pdf_strings)
        tree = etree.ElementTree(root)
        print tree
        
def analyze_alignment_quality(samfile):
    import pysam
    import collections
    import numpy as np
        
    # measures insert size
    isize_hist = collections.defaultdict(lambda: 0)     
    # measures mismatch quality
    mm = np.zeros((2, AlignmentParams.num_mismatches), dtype=np.int)
    read_mm = None
    
    debug_every = 100000
    stop_after = 2000000
    samfhd = pysam.Samfile(samfile, "r")
    for linenum, read in enumerate(samfhd.fetch()):
        # debug
        if linenum % debug_every == 0:
            print linenum
        if linenum == stop_after:
            break
        
        # only consider mapped reads for now
        if read.is_unmapped:
            continue

        if read.rname != read.mrnm:
            print samfhd.getrname(read.rname), samfhd.getrname(read.mrnm)

                
        # tabulate mismatch information
        mismatches = parse_md_str(read.opt('MD'))
        # SAM converts all reads to the forward strand, so if this
        # was a reverse strand read we should change the orientation
        # so that the mismatches are tabulated with respect to the
        # distance from the beginning of the read
        if read.is_reverse:
            mismatches = [(read.rlen-1-x) for x in mismatches]

        # TODO: handle variable read lengths
        mate = 0 if read.is_read1 else 1
        if read_mm is None:
            read_mm = np.zeros((2, read.rlen), dtype=np.int)
        for m in mismatches:
            read_mm[mate,m] += 1
        mm[mate, len(mismatches)] += 1
    
        # keep track of insert sizes
        if read.is_proper_pair and read.is_read1:
            isize_hist[abs(read.isize)] += 1

    samfhd.close()

    # normalize insert size histogram
    sorted_isize_counts = sorted(isize_hist.items(), key=operator.itemgetter(0))
    isizes = np.array([x[0] for x in sorted_isize_counts], dtype=np.int)
    isize_counts = np.array([x[1] for x in sorted_isize_counts], dtype=np.int)
    
    
    aln_stats = AlignmentStats(AlignmentParams())
    aln_stats.read_mm = read_mm
    aln_stats.mm = mm
    aln_stats.isize_mean
    aln_stats.isize_stdev
    aln_stats.isize_pdf

mm_re = re.compile(r'(([0-9]+)([ACGTN]|\^[ACGTN]+))')
 
def parse_md_str(mm_str):
    '''
    given the MD field from a SAM alignment, return the positions
    of the mismatches along the length of the read 
    
    TODO: only handles mismatches now, not indels
    '''
    matches = mm_re.findall(mm_str)
    offset = 0
    mismatches = []
    for m in matches:
        offset += int(m[1]) + 1
        mismatches.append(offset - 1)
    return mismatches


def profile_mate_inner_dist(samfile):
    '''
    Given a SAM file with completely unique reads mapped to the
    transcriptome, produce a histogram of insert sizes and overlay
    with a normal distribution
    '''
    import numpy as np
    import matplotlib.mlab as mlab
    import matplotlib.pyplot as plt
    import pysam
    import array
    import collections
    import matplotlib.patches as patches
    import matplotlib.path as path

    debug_every = 100000
    isize_data = array.array('I')    
    isize_hist = collections.defaultdict(lambda: 0)     
    samfhd = pysam.Samfile(samfile, "r")
    for linenum, read in enumerate(samfhd.fetch()):
        if linenum % debug_every == 0:
            print linenum
        if read.is_proper_pair and read.is_read1 and abs(read.isize < 500):
            isize_hist[abs(read.isize)] += 1
            isize_data.append(abs(read.isize))
    samfhd.close()

#    fig = plt.figure()
#    ax = fig.add_subplot(111)
#
#    # histogram our data with numpy
#    data = np.random.randn(1000)
#    n, bins = np.histogram(data, 50)
#
#    # get the corners of the rectangles for the histogram
#    left = np.array(bins[:-1])
#    right = np.array(bins[1:])
#    bottom = np.zeros(len(left))
#    top = bottom + n
#
#    # we need a (numrects x numsides x 2) numpy array for the path helper
#    # function to build a compound path
#    XY = np.array([[left,left,right,right], [bottom,top,top,bottom]]).T
#    # get the Path object
#    barpath = path.Path.make_compound_path_from_polys(XY)
#    # make a patch out of it
#    patch = patches.PathPatch(barpath, facecolor='blue', edgecolor='gray', alpha=0.8)
#    ax.add_patch(patch)
#    # update the view limits
#    ax.set_xlim(left[0], right[-1])
#    ax.set_ylim(bottom.min(), top.max())
#    plt.show()

    isize_array = np.array(isize_data, dtype=np.int)
    mu = np.average(isize_array)
    sigma = np.std(isize_array)
    median = np.median(isize_array)
    
    #n, bins, patches = plt.hist(isize_array, nbins, normed=1, facecolor='green', alpha=0.75)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    # the histogram of the data
    mybins = np.arange(500)
    n, bins, patches = ax.hist(isize_array, bins=mybins, normed=1, facecolor='green', alpha=0.75)

    bincenters = 0.5*(bins[1:]+bins[:-1])
    # add a 'best fit' line for the normal PDF
    y = mlab.normpdf( bincenters, mu, sigma)
    l = ax.plot(bincenters, y, 'r--', linewidth=1)

    #ax.set_xlim(xmax=500)
    plt.xlabel('Insert Size')
    plt.ylabel('Probability')
    plt.title(r'$\mathrm{Histogram\ of\ ISIZE:}\ \mu=%f,\ \sigma=%f median=%d$' % (mu, sigma, median))    
    #plt.axis([40, 160, 0, 0.03])
    plt.grid(True)

    plt.show()
    
    for k in sorted(isize_hist.keys()):
        print isize_hist[k]


def run_paired_end(mate1_fastq_file, mate2_fastq_file,
                   output_file="bowtie_pe.sam",
                   quals="solexa1.3-quals",
                   processors=2, 
                   genome="hg18_knownGene"):    
    app = Bowtie()
    # options for paired-end alignment
    app.addOptions({'-q': None,
                    '--%s' % (quals): None,
                    '-p': processors,
                    '-v': 2,
                    '-k': 1,
                    '-m': 1,
                    '--best': None,
                    '--sam': None,
                    '-I': 0,
                    '-X': 500000})
    app.args = [app.get_index(genome),
                "-1 %s" % mate1_fastq_file,
                "-2 %s" % mate2_fastq_file, 
                output_file]
    # run the program
    print app.getCommandLine()
    return(app())


from veggie.gene.gene import parse_gene_pred
        
if __name__ == '__main__':
    #run_paired_end(sys.argv[1], sys.argv[2])
    #profile_mate_inner_dist(sys.argv[1])
    analyze_alignment_quality(sys.argv[1])
    sys.exit(0)
    import os
    import random
    
    ucsc_names = {}
    for g in parse_gene_pred(open('/lab/mkiyer/refdb/ucsc_tables/knownGene.txt')):
        ucsc_names[g[0]] = g[2] - g[1]
    n = len(ucsc_names)
    ucsc_keys = ucsc_names.keys()
    
    for i in xrange(1000):
        key = random.randint(0, n)
        print i
        os.system('samtools view sorted.bam %s > /dev/null' % ucsc_keys[key])
        


